import torch
import torch.nn as nn
import math

def autopad(k, p=None, d=1):  # kernel, padding, dilation
    """自动计算padding大小"""
    if d > 1:
        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
    return p

class Conv(nn.Module):
    """标准卷积模块"""
    default_act = nn.SiLU()  # default activation

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

class Bottleneck(nn.Module):
    """标准瓶颈结构"""
    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, k[0], 1)
        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

class C2f(nn.Module):
    """CSP Bottleneck with 2 convolutions and fusion"""
    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
        super().__init__()
        self.c = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

    def forward(self, x):
        y = list(self.cv1(x).chunk(2, 1))
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

class SPPF(nn.Module):
    """Spatial Pyramid Pooling - Fast (SPPF) layer"""
    def __init__(self, c1, c2, k=5):
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        y1 = self.m(x)
        y2 = self.m(y1)
        return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))

def make_anchors(feats, strides, grid_cell_offset=0.5):
    """生成锚点"""
    anchor_points, stride_tensor = [], []
    for i, stride in enumerate(strides):
        _, _, h, w = feats[i].shape
        sx = torch.arange(w) + grid_cell_offset  # shift x
        sy = torch.arange(h) + grid_cell_offset  # shift y
        sy, sx = torch.meshgrid(sy, sx, indexing='ij')
        anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2))
        stride_tensor.append(torch.full((h * w, 1), stride))
    return torch.cat(anchor_points), torch.cat(stride_tensor)

class YOLOv8(nn.Module):
    """YOLOv8 object detection model"""
    def __init__(self, nc=80, ch=3, model_size='n'):  # number of classes, channels
        super().__init__()
        # 根据模型大小选择深度和宽度因子
        depth_dict = {'n': 0.33, 's': 0.33, 'm': 0.67}
        width_dict = {'n': 0.25, 's': 0.50, 'm': 0.75}
        depth = depth_dict[model_size]
        width = width_dict[model_size]
        
        # 基础通道数
        channels = {
            'n': (16, 32, 64, 128, 256),
            's': (32, 64, 128, 256, 512),
            'm': (48, 96, 192, 384, 768),
        }[model_size]

        # Backbone
        self.backbone = nn.ModuleList([
            # P1/2
            Conv(ch, channels[0], 3, 2),
            # P2/4
            Conv(channels[0], channels[1], 3, 2),
            C2f(channels[1], channels[1], round(3 * depth)),
            # P3/8
            Conv(channels[1], channels[2], 3, 2),
            C2f(channels[2], channels[2], round(6 * depth)),
            # P4/16
            Conv(channels[2], channels[3], 3, 2),
            C2f(channels[3], channels[3], round(6 * depth)),
            # P5/32
            Conv(channels[3], channels[4], 3, 2),
            C2f(channels[4], channels[4], round(3 * depth)),
            SPPF(channels[4], channels[4], k=5)
        ])

        # Neck
        self.neck = nn.ModuleList([
            # 上采样路径
            Conv(channels[4], channels[3], 1, 1),
            nn.Upsample(None, 2, 'nearest'),
            C2f(channels[3] * 2, channels[3], round(3 * depth)),
            
            Conv(channels[3], channels[2], 1, 1),
            nn.Upsample(None, 2, 'nearest'),
            C2f(channels[2] * 2, channels[2], round(3 * depth)),
            
            # 下采样路径
            Conv(channels[2], channels[2], 3, 2),
            C2f(channels[2] * 2, channels[3], round(3 * depth)),
            
            Conv(channels[3], channels[3], 3, 2),
            C2f(channels[3] * 2, channels[4], round(3 * depth))
        ])

        # Head
        self.head = nn.ModuleList([
            nn.Conv2d(channels[2], nc + 4, 1),  # P3/8
            nn.Conv2d(channels[3], nc + 4, 1),  # P4/16
            nn.Conv2d(channels[4], nc + 4, 1)   # P5/32
        ])

        self.nc = nc
        self.reg_max = 16  # DFL channels
        self.stride = torch.tensor([8., 16., 32.])
        
        # 初始化权重
        self.initialize_weights()

    def initialize_weights(self):
        """Initialize model weights"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        # Backbone
        features = []
        for i, m in enumerate(self.backbone):
            if i in [2, 4, 6, 9]:  # 保存特征图
                features.append(x)
            x = m(x)
        features.append(x)
        
        # Neck
        fpn_features = []
        x = features[-1]
        for i, m in enumerate(self.neck):
            if i in [0, 3]:  # 上采样前的卷积
                x = m(x)
            elif i in [1, 4]:  # 上采样
                x = m(x)
                x = torch.cat([x, features[3-i//3]], 1)
            elif i in [6, 8]:  # 下采样
                x = m(x)
                x = torch.cat([x, fpn_features[i//3-2]], 1)
            else:  # C2f
                x = m(x)
                fpn_features.append(x)
        
        # Head
        outputs = []
        for i, m in enumerate(self.head):
            outputs.append(m(fpn_features[i]))
            
        return outputs

    def _make_grid(self, nx=20, ny=20):
        """生成网格点"""
        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)], indexing='ij')
        return torch.stack((xv, yv), 2).view(-1, 2).float()

def build_model(name='yolov8n', nc=80):
    """构建YOLOv8模型"""
    size_map = {'yolov8n': 'n', 'yolov8s': 's', 'yolov8m': 'm'}
    if name not in size_map:
        raise ValueError(f"Unsupported model name: {name}")
    return YOLOv8(nc=nc, model_size=size_map[name])
